import sys

import cv2
from PIL import Image,ImageDraw,ImageFont
import numpy as np

sys.path.append("../")
from plate_detect.plate_detect import PlateDetect

class VideoProcessing:

    # 初始化
    def __init__(self, path=r"../plate_detect/ultralytics/weights/best.pt"):
        # 模型加载
        self.plate_engine = PlateDetect(path)    

    # 视频加载
    def Video(self, path):
        return cv2.VideoCapture(path)
    
    def run(self, path):
        cv2.namedWindow("win", 0)
        cap = self.Video(path)
        while(cap.isOpened()):
            try:
                ret, frame = cap.read()
                image, total, dish_list = self.plate_engine.detect_frame(frame)
                pil_img = Image.fromarray(cv2.cvtColor(image,cv2.COLOR_BGR2RGB))
                draw = ImageDraw.Draw(pil_img)
                font = ImageFont.truetype('simhei.ttf', 100, encoding="utf-8")
                draw.text((50,200), f"总价：{total}", (255,0,0), font)
                image = cv2.cvtColor(np.array(pil_img),cv2.COLOR_RGB2BGR)
            except Exception as e:
                break

            # 获取成功显示图像
            if ret == True:
                cv2.imshow("win", image)
            # 每一帧间隔为25ms
            if cv2.waitKey(25) & 0xFF == ord('q'):
                break

        # 释放视频对象
        cap.release()
        cv2.destroyAllWindows()

if __name__ == "__main__":
    video_processing = VideoProcessing()
    video_processing.run("../plate_detect/ultralytics/data/test.mp4")